from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("/data1/jzh/models/Llama-2-7B-fp16")

# 设置 Llama-2 官方聊天模板
tokenizer.chat_template = """
<s>[INST] {% if messages[0]["role"] == "system" %}
<<SYS>>
{{ messages[0]["content"] }}
<</SYS>>
{% endif %}
{% if messages[0]["role"] == "user" %}{{ messages[0]["content"] }}{% endif %}
[/INST]"""

print(tokenizer.chat_template)